from common import setup_logger, read_task_info_file, upsert_task_info, download_file, TASK_STATUS, image_to_base64
from klingai import KlingAI


class TaskManager:
    def __init__(self):
        self.log = setup_logger('task_manager')

    def update_all_task_status(self):
        data = read_task_info_file()
        kling_ai = KlingAI()

        new_array = []
        for old_task_data in data:
            task_id = old_task_data['task_id']
            self.log.info(f'task_id: {task_id}')
            if 'success' in old_task_data and old_task_data['success']:
                self.log.info('status succeed')
            else:
                task_status, videos_array = kling_ai.query_task_status(task_id)

                if task_status == TASK_STATUS.SUCCESS:
                    for video in videos_array:
                        id = video['id']
                        url = video['url']
                        filename = f'{id}.mp4'
                        download_file(url, filename)

                    task = {
                        'task_id': task_id,
                        'success': True,
                        'videos': videos_array
                    }
                    upsert_task_info(task)

    def demo_image_to_video(self):
        """
        简单的官方示例，示例中仅支持1.0的模型
        :return:
        """
        kling_ai = KlingAI()

        model_name = 'kling-v1'  # kling-v1, kling-v1-5
        mode = 'pro'  # 枚举值：std(标准模式（标准），基础模式，性价比高)，pro(专家模式（高品质），高表现模式，生成视频质量更佳)
        duration = 5  # 单位：s秒，枚举5，10,
        image = 'https://h2.inkwai.com/bs2/upload-ylab-stunt/se/ai_portal_queue_mmu_image_upscale_aiweb/3214b798-e1b4-4b00-b7af-72b5b0417420_raw_image_0.jpg'
        prompt = '宇航员站起身走了'
        cfg_scale = 0.5
        negative_prompt = None
        static_mask = 'https://h2.inkwai.com/bs2/upload-ylab-stunt/ai_portal/1732888177/cOLNrShrSO/static_mask.png'
        dynamic_masks = [{
            "mask": "https://h2.inkwai.com/bs2/upload-ylab-stunt/ai_portal/1732888130/WU8spl23dA/dynamic_mask_1.png",
            "trajectories": [
                {"x": 279, "y": 219}, {"x": 417, "y": 65}
            ]
        }]
        create_task_response = kling_ai.image_to_video(model_name, mode, duration, image, prompt, cfg_scale, negative_prompt, static_mask, dynamic_masks)

        if create_task_response['code'] == 0:
            data_node = create_task_response['data']
            task_id = data_node['task_id']
            task_status = data_node['task_status']
            self.log.info(f'create new task, task_id: {task_id}, status: {task_status}')

            upsert_task_info({
                'task_id': task_id
            })
        else:
            self.log.error(create_task_response)

    def simple_image_to_video(self,
                              local_image_path: str,
                              prompt: str,
                              negative_prompt: str,
                              cfg_scale: float = 0.5,
                              duration: int = 5):
        """

        :param local_image_path: 图片本地路径
        :param prompt: 正向提示词
        :param negative_prompt: 负向提示词
        :param cfg_scale: 创意想象力/创意相关性
        :param duration: 生成视频时长，5或10
        :return:
        """
        image_base64 = image_to_base64(local_image_path)

        kling_ai = KlingAI()
        model_name = 'kling-v1-5'
        mode = 'pro'
        static_mask = None
        dynamic_masks = None
        create_task_response = kling_ai.image_to_video(model_name, mode, duration, image_base64, prompt, cfg_scale, negative_prompt, static_mask, dynamic_masks)

        if create_task_response['code'] == 0:
            data_node = create_task_response['data']
            task_id = data_node['task_id']
            task_status = data_node['task_status']
            self.log.info(f'create new task, task_id: {task_id}, status: {task_status}')

            upsert_task_info({
                'task_id': task_id,
                'local_image_path': local_image_path
            })
        else:
            self.log.error(create_task_response)

